# 修改后的 fake_coin_puzzle.py
import random
from typing import Annotated

from app.puzzle.schema import _IE


def get_matrix_height(n: int):
    """
    计算区分 n 个硬币的最小称量次数
    """
    if n <= 2:
        raise ValueError("硬币数必须大于2")

    h = 0
    while True:
        h += 1
        # 计算可用向量数：(3^h - 1) // 2
        available_vectors = (3**h - 1) // 2
        if available_vectors > n:
            return h


def adjust_matrix_to_balance(combos: list[list]) -> list[list]:
    """
    将矩阵调整为每行和为0
    """
    combos_list: list = list(combos)
    # 添加尝试次数，避免死循环
    max_try = len(combos)
    count = 0
    while True:
        adjust_list = []
        adjust_value_list = []
        for i in range(len(combos_list[0])):
            # 获取偏差，并归一化
            sum_row = sum([x[i] for x in combos_list])
            adjust_value_list.append(sum_row)
            if sum_row > 0:
                adjust_list.append(1)
            elif sum_row < 0:
                adjust_list.append(-1)
            else:
                adjust_list.append(0)
        # 如果所有行和为0，则返回
        if all(x == 0 for x in adjust_list):
            break
        try:
            # 根据偏差找到可调整的列，并把尝试次数置为0
            index = combos_list.index(adjust_list)
            combos_list[index] = [x * -1 for x in combos_list[index]]
            count = 0
        except ValueError:
            # 如果找不到可调整的列，则找到偏差最大的行，随机取可以调整的列，尝试次数+1
            adjust_list_abs = [abs(x) for x in adjust_value_list]
            max_sum_abs = max(adjust_list_abs)
            index_max = adjust_list_abs.index(max_sum_abs)
            value_max = adjust_list[index_max]
            index_list = []
            for i in range(len(combos_list)):
                if combos_list[i][index_max] * value_max > 0:
                    index_list.append(i)
            # 如果偏差很大，一次调整多列
            if max_sum_abs > 10:
                indexes = random.sample(index_list, int(max_sum_abs / 2))
                for i in indexes:
                    combos_list[i] = [x * -1 for x in combos_list[i]]
            # 偏差不大，一次调整1列
            else:
                index = random.choice(index_list)
                combos_list[index] = [x * -1 for x in combos_list[index]]

            # 随机调整需要有尝试上限，避免进入死循环
            count += 1
            if count > max_try:
                return [[]]
    return combos_list


def get_adjust_matrix_to_balance(combos):
    valid = True
    # 每一行的和都要是偶数, 否则对任意列取反，无法保证每行和为0
    for i in range(len(combos[0])):
        if sum([x[i] for x in combos]) % 2 != 0:
            valid = False
            break
    # 确保每行和为偶数， 不满足条件，使用下一个组合
    if not valid:
        return None
    # 通过取反向向量，使每行和为0
    if adjust_arrays := adjust_matrix_to_balance(combos):
        return adjust_arrays
    return None


def select_available_groups_plan_c(original_arrays, n):
    # 选取函数还能优化 TODO
    """
    从所有向量中选出可能调整的向量
    """
    max_attempts = 10000
    for _ in range(max_attempts):
        # 随机选取组合，确保每行和为偶数
        # 注意这里不能使用排列组合按顺序选取，因为前面一部分总是不满足要求的，当数量级较大时，会耗时过长，使用随机选取，概率是一定的，选不到的概率很低
        combos = random.sample(original_arrays, n)
        if adjust_arrays := get_adjust_matrix_to_balance(combos):
            return adjust_arrays
        else:
            continue
    raise ValueError("无法找到合适的组合")


def select_available_groups_plan_b(original_arrays, n):
    """
    从所有向量中选出可能调整的向量
    """
    max_attempts = 10000
    for _ in range(max_attempts):
        # 随机选取组合，确保每行和为偶数
        # 注意这里不能使用排列组合按顺序选取，因为前面一部分总是不满足要求的，当数量级较大时，会耗时过长，使用随机选取，概率是一定的，选不到的概率很低
        group_zero = [x for x in original_arrays if x[0] == 0]
        group_one = [x for x in original_arrays if x[0] != 0]
        group_zero_one = []
        group_zero_other = []
        h = get_matrix_height(n)
        indexes = []
        for i in range(h):
            indexes.append(int((3**i - 1) / 2))
        for i in range(len(group_zero)):
            if i in indexes:
                group_zero_one.append(group_zero[i])
            else:
                group_zero_other.append(group_zero[i])
        if n % 3 == 0:
            zero_count = int(n / 3)
        elif n % 3 == 1:
            zero_count = int((n + 2) / 3)
        else:
            zero_count = int((n - 2) / 3)

        one_count = int(n - zero_count)

        combos = (
            group_zero_one
            + random.sample(group_zero_other, zero_count - len(group_zero_one))
            + [group_one[0]]
            + random.sample(group_one[1:], one_count - 1)
        )

        if adjust_arrays := get_adjust_matrix_to_balance(combos):
            return adjust_arrays
        else:
            continue

    raise ValueError("无法找到合适的组合")


def select_available_groups_plan_a(original_arrays, n):
    # 首先，将original_arrays 分组
    groups = group_original_arrays(original_arrays)
    # 从groups选择一些，使得其count相加正好等于n

    valid_group_combination = find_valid_group_combination(groups, n)
    valid_arrays = []
    for sub_group in valid_group_combination:
        valid_arrays.extend(sub_group["vec_sub"])
    if adjust_arrays := adjust_matrix_to_balance(valid_arrays):
        return adjust_arrays
    raise ValueError("无法找到合适的组合")


def find_valid_group_combination(groups: list[dict], target_count: int) -> list[dict]:
    """
    找出满足条件的组合（只需一组）：
    1. 分组计数之和等于 target_count
    2. first_one_count 字段之和为偶数
    """
    # 按分组大小降序排序（贪心策略）
    groups.sort(key=lambda x: x["count"], reverse=True)

    # 计算分组总大小（用于剪枝）
    total_count = sum(g["count"] for g in groups)

    # 回溯搜索函数
    def backtrack(
        start: int, curr_sum: int, curr_parity: int, path: list[dict]
    ) -> list[dict] | None:
        # 找到满足条件的解
        if curr_sum == target_count and curr_parity % 2 == 0:
            return path.copy()

        # 剪枝条件
        if curr_sum > target_count:
            return None
        if curr_sum + total_count < target_count:  # 剩余分组总大小不足
            return None

        # 尝试后续分组
        for i in range(start, len(groups)):
            group = groups[i]
            # 选择当前分组
            new_sum = curr_sum + group["count"]
            new_parity = curr_parity + group["first_one_count"]

            # 递归搜索
            path.append(group)
            _result = backtrack(i + 1, new_sum, new_parity, path)
            if _result:
                return _result
            path.pop()  # 回溯

        return None

    # 开始搜索
    result = backtrack(0, 0, 0, [])
    if result:
        return result  # 保持返回类型一致

    raise ValueError(
        f"无法找到总和为{target_count}且first_one_count总和为偶数的分组组合"
    )


def group_original_arrays(original_arrays: list[list[_IE]]):
    """
    将 original_arrays 分组，返回保存了第一行1的数量（根据对称性，每行的0的数量都相等，保证了每行的奇偶性和第一行相同）
    """
    n = len(original_arrays[0])

    # 先根据每一列0的数量
    group_step_one: list[list] = [[] for _ in range(n)]
    for vec in original_arrays:
        # 计算0的数量
        zero_count = sum(1 for x in vec if x == 0)
        group_step_one[zero_count].append(vec)

    # 根据每一列-1的数量
    group_step_two = []
    for group in group_step_one:
        group_step_two_sub: list[list] = [[] for _ in range(n)]
        for vec in group:
            # 计算-1的数量
            minus_one_count = sum(1 for x in vec if x == -1)
            group_step_two_sub[minus_one_count].append(vec)
        for vec_sub in group_step_two_sub:
            if not vec_sub:
                continue
            # 计算第一行1的数量，用于分辨奇偶性
            first_one_count = sum(1 for x in vec_sub if x[0] == 1)
            group_step_two.append(
                {
                    "first_one_count": first_one_count,
                    "vec_sub": vec_sub,
                    "count": len(vec_sub),
                }
            )

    return group_step_two


# 修改后的矩阵构建函数（不依赖NumPy）
def build_measurement_matrix(n: int):
    # 获取最小称量次数
    h = get_matrix_height(n)

    # 获取向量总数
    total_vectors = 3**h

    # 第一步：生成所有向量， 排除零向量
    matrix = []
    for i in range(1, total_vectors):
        vector: list[Annotated[int, _IE]] = []
        num = i
        # 转换为h位三进制（高位在前）
        for _ in range(h):
            num, r = divmod(num, 3)
            r_value = -1 if r == 2 else r
            vector.insert(0, r_value)  # 插入开头保证高位在前
        matrix.append(vector)

    # 第二步，分组，相反向量只保留一个
    original_arrays = []
    used = set()
    for vec in matrix:
        vec_tuple = tuple(vec)
        if vec_tuple not in used:
            opposite = tuple(-x for x in vec)
            original_arrays.append(list(vec))
            used.add(vec_tuple)
            used.add(opposite)

    # 第三步，取出n组
    try:
        result = select_available_groups_plan_a(original_arrays, n)
        return result
    except ValueError:
        pass

    try:
        result = select_available_groups_plan_b(original_arrays, n)
        return result
    except ValueError:
        pass

    try:
        result = select_available_groups_plan_c(original_arrays, n)
        return result
    except ValueError:
        pass

    raise ValueError("无法找到合适的组合")


def _format_measurement_description(coin_indices: list[str], side: str) -> str:
    """
    格式化单次称量的描述

    Args:
        coin_indices: 硬币编号列表
        side: 天平的一侧 ("左边" 或 "右边")

    Returns:
        格式化后的描述字符串
    """
    if coin_indices:
        return f"{side}放 {', '.join(coin_indices)}号硬币"
    return ""


def convert_to_plain_language(matrix):
    # 将matrix转为通俗易懂的语言
    h = len(matrix[0])
    result = f"一共{len(matrix)}枚硬币，你需要进行{h}次称量: \n"
    for i in range(h):
        left = []
        right = []

        # 收集当前称量的硬币分配
        for coin_idx in range(len(matrix)):
            val = matrix[coin_idx][i]
            if val == 1:
                left.append(str(coin_idx + 1))  # 硬币编号从1开始
            elif val == -1:
                right.append(str(coin_idx + 1))

        # 构建称量描述
        result += f"第{i + 1}次称量："
        left_desc = _format_measurement_description(left, "左边")
        right_desc = _format_measurement_description(right, "右边")

        if left_desc and right_desc:
            result += f"{left_desc}，{right_desc}"
        elif left_desc or right_desc:
            result += left_desc + right_desc
        else:
            result += "两边都不放硬币"
        result += "\n"
    return result


def fake_coin_puzzle_get_method(n: int) -> tuple[bool, dict]:
    try:
        matrix = build_measurement_matrix(n)
        return True, {"data": matrix, "message": convert_to_plain_language(matrix)}
    except Exception as e:
        return False, {"message": str(e)}


def fake_coin_puzzle_get_result(
    matrix: list[list[_IE]], result_list: list[_IE]
) -> tuple[bool, dict]:
    try:
        # 如果矩阵测量次数和result_list不相等，返回错误
        if len(matrix[0]) != len(result_list):
            return False, {"message": "测量次数和结果数量不相等"}
        explanation = convert_to_plain_language(matrix) + "\n称量结果是：\n"
        for i in range(len(result_list)):
            if result_list[i] > 0:
                explanation += f"第{i + 1}次称量：左边较重 \n"
            elif result_list[i] < 0:
                explanation += f"第{i + 1}次称量：右边较重 \n"
            else:
                explanation += f"第{i + 1}次称量：天平平衡 \n"

        # 如果结果向量是零向量
        if all(x == 0 for x in result_list):
            return True, {"message": f"{explanation}\n没有假币", "data": 0}
        # 从矩阵中找到满足条件的向量， 满足的条件为元素等于result_list ,或者是他的反向向量
        for i in range(len(matrix)):
            if matrix[i] == result_list:
                return True, {
                    "message": f"{explanation}\n第{i + 1}个硬币较重",
                    "data": i + 1,
                }
            if matrix[i] == [-x for x in result_list]:
                return True, {
                    "message": f"{explanation}\n第{i + 1}个硬币较轻",
                    "data": -(i + 1),
                }
        return False, {"message": f"{explanation}\n没有找到对应的硬币"}
    except Exception as e:
        return False, {"message": str(e)}


# 测试示例
# if __name__ == "__main__":
#     n = 12
#     # start = time.time()
#     # matrix = build_measurement_matrix(n)
#     # print(time.time() - start)
#     #
#     ok, method = fake_coin_puzzle_get_method(n)
#     if ok:
#         data = method.get("data")
#         message = method.get("message")
#         print(data)
#         result_list = []
#         for i in range(len(data[0])):
#             # 使用input获取输入
#             result = int(
#                 input(
#                     f"请输入第{i + 1}次称量结果（天平左边重为1，右边重为-1，平为0）："
#                 )
#             )
#             result_list.append(int(str(result).strip()))
#         ok, result = fake_coin_puzzle_get_result(data, result_list)
#         print(result.get("message"))
#     else:
#         print(method.get("message"))
